#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import math

class FrenetPlanner:
    def __init__(self, reference_path):
        """初始化Frenet规划器
        
        参数:
            reference_path: 参考路径点的列表，每个点是(x, y)元组
        """
        self.reference_path = np.array(reference_path)
        self.s_list = self._calculate_s_values()
        
    def _calculate_s_values(self):
        """计算参考路径上每个点的累积距离(s值)"""
        s = [0.0]
        for i in range(1, len(self.reference_path)):
            dx = self.reference_path[i, 0] - self.reference_path[i-1, 0]
            dy = self.reference_path[i, 1] - self.reference_path[i-1, 1]
            s.append(s[-1] + math.hypot(dx, dy))
        return np.array(s)
    
    def _find_closest_point(self, x, y):
        """找到参考路径上离给定点最近的点"""
        min_dist = float('inf')
        closest_idx = 0
        
        for i in range(len(self.reference_path)):
            dx = x - self.reference_path[i, 0]
            dy = y - self.reference_path[i, 1]
            dist = dx*dx + dy*dy  # 用平方避免开方，提高效率
            
            if dist < min_dist:
                min_dist = dist
                closest_idx = i
                
        # 检查下一个点是否更近（线性插值）
        if closest_idx < len(self.reference_path) - 1:
            a = self.reference_path[closest_idx]
            b = self.reference_path[closest_idx + 1]
            
            # 向量AB
            ab = b - a
            # 向量AP
            ap = np.array([x, y]) - a
            
            # 计算投影比例
            proj = np.dot(ap, ab) / (np.dot(ab, ab) + 1e-8)
            
            if 0 <= proj <= 1:
                # 计算投影点
                proj_point = a + proj * ab
                dx_proj = x - proj_point[0]
                dy_proj = y - proj_point[1]
                proj_dist = dx_proj*dx_proj + dy_proj*dy_proj
                
                if proj_dist < min_dist:
                    s = self.s_list[closest_idx] + proj * np.linalg.norm(ab)
                    return closest_idx, proj, s, proj_point
        
        return closest_idx, 0.0, self.s_list[closest_idx], self.reference_path[closest_idx]
    
    def cartesian_to_frenet(self, x, y, theta):
        """将笛卡尔坐标转换为Frenet坐标
        
        参数:
            x, y: 笛卡尔坐标
            theta: 车辆航向角（弧度）
            
        返回:
            s: 沿参考路径的距离
            d: 垂直于参考路径的距离
            ds: 沿参考路径方向的速度分量
            dd: 垂直于参考路径方向的速度分量
        """
        closest_idx, proj, s, closest_point = self._find_closest_point(x, y)
        
        # 计算参考路径在最近点的方向
        if closest_idx < len(self.reference_path) - 1:
            dx_ref = self.reference_path[closest_idx + 1, 0] - self.reference_path[closest_idx, 0]
            dy_ref = self.reference_path[closest_idx + 1, 1] - self.reference_path[closest_idx, 1]
        else:
            dx_ref = self.reference_path[closest_idx, 0] - self.reference_path[closest_idx - 1, 0]
            dy_ref = self.reference_path[closest_idx, 1] - self.reference_path[closest_idx - 1, 1]
            
        theta_ref = math.atan2(dy_ref, dx_ref)
        
        # 计算d值（垂直距离）
        dx = x - closest_point[0]
        dy = y - closest_point[1]
        d = math.hypot(dx, dy)
        
        # 确定d的符号（左侧为正，右侧为负）
        cross = dx * math.sin(theta_ref) - dy * math.cos(theta_ref)
        if cross < 0:
            d = -d
        
        # 简化的速度分量计算（假设速度为单位速度）
        ds = math.cos(theta - theta_ref)
        dd = math.sin(theta - theta_ref)
        
        return s, d, ds, dd
    
    def frenet_to_cartesian(self, s, d):
        """将Frenet坐标转换为笛卡尔坐标
        
        参数:
            s: 沿参考路径的距离
            d: 垂直于参考路径的距离
            
        返回:
            x, y: 笛卡尔坐标
            theta: 该点的航向角
        """
        # 找到s对应的参考路径点
        idx = np.searchsorted(self.s_list, s) - 1
        if idx < 0:
            idx = 0
        if idx >= len(self.s_list) - 1:
            idx = len(self.s_list) - 2
            
        # 计算插值比例
        s_segment = self.s_list[idx + 1] - self.s_list[idx]
        if s_segment < 1e-8:
            ratio = 0.0
        else:
            ratio = (s - self.s_list[idx]) / s_segment
            
        # 插值得到参考点
        x_ref = self.reference_path[idx, 0] + ratio * (self.reference_path[idx + 1, 0] - self.reference_path[idx, 0])
        y_ref = self.reference_path[idx, 1] + ratio * (self.reference_path[idx + 1, 1] - self.reference_path[idx, 1])
        
        # 计算参考路径方向
        dx_ref = self.reference_path[idx + 1, 0] - self.reference_path[idx, 0]
        dy_ref = self.reference_path[idx + 1, 1] - self.reference_path[idx, 1]
        theta_ref = math.atan2(dy_ref, dx_ref)
        
        # 计算垂直方向
        theta_perp = theta_ref + math.pi / 2
        
        # 计算笛卡尔坐标
        x = x_ref + d * math.cos(theta_perp)
        y = y_ref + d * math.sin(theta_perp)
        
        return x, y, theta_ref
    
    def plan_trajectory(self, start_s, start_d, end_s, end_d, num_points=100):
        """规划从起点到终点的轨迹
        
        参数:
            start_s, start_d: 起点的Frenet坐标
            end_s, end_d: 终点的Frenet坐标
            num_points: 轨迹点数量
            
        返回:
            轨迹点的列表，每个点是(s, d)元组
        """
        # 简单的五次多项式规划
        s = np.linspace(start_s, end_s, num_points)
        s_norm = (s - start_s) / (end_s - start_s + 1e-8)
        
        # 五次多项式参数
        a0 = start_d
        a1 = 0  # 初始横向速度为0
        a2 = 0  # 初始横向加速度为0
        a3 = 10 * (end_d - start_d)
        a4 = -15 * (end_d - start_d)
        a5 = 6 * (end_d - start_d)
        
        d = a0 + a1 * s_norm + a2 * s_norm**2 + a3 * s_norm**3 + a4 * s_norm**4 + a5 * s_norm**5
        
        return np.column_stack((s, d))

def create_animation(planner, trajectory):
    """创建轨迹动画
    
    参数:
        planner: FrenetPlanner实例
        trajectory: 规划的轨迹点列表
    """
    # 转换轨迹到笛卡尔坐标
    cartesian_traj = [planner.frenet_to_cartesian(s, d) for s, d in trajectory]
    cartesian_traj = np.array(cartesian_traj)
    
    # 设置图形
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.set_xlabel('X Position')
    ax.set_ylabel('Y Position')
    ax.set_title('Frenet Planner Trajectory')
    ax.grid(True)
    
    # 绘制参考路径
    ax.plot(planner.reference_path[:, 0], planner.reference_path[:, 1], 'b--', label='Reference Path')
    
    # 初始化轨迹线和车辆标记
    traj_line, = ax.plot([], [], 'r-', label='Planned Trajectory')
    vehicle_marker, = ax.plot([], [], 'go', markersize=10, label='Vehicle')
    
    # 设置坐标轴范围
    all_x = np.concatenate([planner.reference_path[:, 0], cartesian_traj[:, 0]])
    all_y = np.concatenate([planner.reference_path[:, 1], cartesian_traj[:, 1]])
    ax.set_xlim(min(all_x) - 1, max(all_x) + 1)
    ax.set_ylim(min(all_y) - 1, max(all_y) + 1)
    ax.legend()
    
    def init():
        traj_line.set_data([], [])
        vehicle_marker.set_data([], [])
        return traj_line, vehicle_marker
    
    def update(frame):
        # 更新轨迹线
        traj_line.set_data(cartesian_traj[:frame+1, 0], cartesian_traj[:frame+1, 1])
        # 更新车辆位置
        vehicle_marker.set_data(cartesian_traj[frame, 0], cartesian_traj[frame, 1])
        return traj_line, vehicle_marker
    
    # 创建动画
    anim = FuncAnimation(fig, update, frames=len(cartesian_traj),
                         init_func=init, blit=True, interval=50)
    
    plt.show()
    
    return anim

def main():
    # 定义参考路径点（可以修改这些点来改变参考路径）
    reference_path = [
        (0, 0),    # 起点
        (5, 0.1),    # 直线路径
        (10, 0.3),   # 开始右转
        (15, 0.6),   # 继续右转
        (20, 0.8),  # 直线路径
        (25, 0.6),  # 直线路径
        (30, 0.4),   # 开始左转
        (35, 0.2),   # 继续左转
        (40, 0)    # 终点
    ]
    wx = [44.90, 44.89, 44.89, 45.07, 45.35, 46.08, 46.83, 47.81, 48.31]
    wy = [-33.86, -32.94, -31.97, -31.17, -30.9, -30.54, -30.3, -30.14, -30.1]
    reference_path = [(x, y) for x, y in zip(wx, wy)]
    
    # 创建Frenet规划器
    planner = FrenetPlanner(reference_path)
    
    # 规划轨迹（从参考路径起点稍微偏移的位置到终点稍微偏移的位置）
    start_s = 0.0
    start_d = 0.2  # 起点在参考路径右侧1单位
    end_s = planner.s_list[-1]
    end_d = 0.0   # 终点在参考路径左侧1单位
    
    trajectory = planner.plan_trajectory(start_s, start_d, end_s, end_d, 200)
    
    # 创建并显示动画
    create_animation(planner, trajectory)

if __name__ == "__main__":
    main()
